from ultralytics import YOLO
from ultralytics.models.yolo.segment import SegmentationPredictor
import cv2
import numpy as np
import os
import matplotlib.pyplot as plt
from rknn.api import RKNN
import torch

def visualize_and_save_masks(results, original_img, output_dir="output"):
    """可视化并保存分割掩码结果"""
    os.makedirs(output_dir, exist_ok=True)
    
    orig_path = os.path.join(output_dir, "original.jpg")
    cv2.imwrite(orig_path, original_img)
    print(f"原始图像已保存至: {orig_path}")
    
    if not results or len(results) == 0:
        print("无检测结果可保存")
        return
    
    pred_path = os.path.join(output_dir, "prediction.jpg")
    results[0].save(filename=pred_path)
    print(f"预测结果图像已保存至: {pred_path}")
    
    if results[0].masks is not None:
        masks = results[0].masks.data.cpu().numpy()
        boxes = results[0].boxes.xyxy.cpu().numpy()
        
        for i, (mask, box) in enumerate(zip(masks, boxes)):
            binary_mask = (mask > 0).astype(np.uint8) * 255
            
            if binary_mask.shape != original_img.shape[:2]:
                binary_mask = cv2.resize(binary_mask, 
                                        (original_img.shape[1], original_img.shape[0]), 
                                        interpolation=cv2.INTER_NEAREST)
            
            mask_path = os.path.join(output_dir, f"mask_{i}.png")
            cv2.imwrite(mask_path, binary_mask)
            print(f"实例 {i} 的分割掩码已保存至: {mask_path}")
    else:
        print("未检测到分割掩码")

def init_rknn_enhanced(rknnModel, id=1):
    """增强的RKNN初始化函数"""
    print(f"\n=== RKNN模型初始化 ===")
    print(f"模型路径: {rknnModel}")
    
    if not os.path.exists(rknnModel):
        print(f"错误: 模型文件不存在: {rknnModel}")
        return None
    
    model_size = os.path.getsize(rknnModel) / (1024 * 1024)
    print(f"模型文件大小: {model_size:.2f} MB")
    
    rknn = RKNN(verbose=True)
    
    print("加载RKNN模型...")
    ret = rknn.load_rknn(rknnModel)
    if ret != 0:
        print(f"模型加载失败，错误代码: {ret}")
        return None
    print("✓ 模型加载成功")
    
    print("初始化运行时环境...")
    core_configs = {
        0: (RKNN.NPU_CORE_0, "NPU_CORE_0"),
        1: (RKNN.NPU_CORE_1, "NPU_CORE_1"),
        2: (RKNN.NPU_CORE_2, "NPU_CORE_2"),
        -1: (RKNN.NPU_CORE_0_1_2, "NPU_CORE_0_1_2")
    }
    
    if id in core_configs:
        core_mask, core_name = core_configs[id]
        ret = rknn.init_runtime(core_mask=core_mask, target='rk3588')
        print(f"✓ 模型加载到 {core_name}")
    else:
        ret = rknn.init_runtime(target='rk3588')
        print("✓ 模型加载到默认NPU")
        
    if ret != 0:
        print(f"运行时环境初始化失败，错误代码: {ret}")
        return None
    
    print("✓ 运行时环境初始化成功")
    return rknn

def validate_and_prepare_input(pre_result):
    """验证并准备输入数据"""
    print(f"\n=== 输入数据准备 ===")
    
    # 转换为numpy
    input_data = pre_result.numpy()
    
    print(f"输入形状: {input_data.shape}")
    print(f"输入数据类型: {input_data.dtype}")
    print(f"输入数据范围: [{input_data.min():.6f}, {input_data.max():.6f}]")
    print(f"输入数据统计: 均值={input_data.mean():.6f}, 标准差={input_data.std():.6f}")
    
    # 检查输入格式
    if len(input_data.shape) != 4:
        print(f"错误: 输入形状不正确，期望4维，实际{len(input_data.shape)}维")
        return None
    
    if input_data.shape[1] != 3:
        print(f"错误: 通道数不正确，期望3通道，实际{input_data.shape[1]}通道")
        return None
    
    # 确保数据类型为float32
    if input_data.dtype != np.float32:
        print(f"转换数据类型: {input_data.dtype} -> float32")
        input_data = input_data.astype(np.float32)
    
    # 检查数据有效性
    if np.isnan(input_data).any():
        print("错误: 输入数据包含NaN值")
        return None
    
    if np.isinf(input_data).any():
        print("错误: 输入数据包含无穷值")
        return None
    
    # 保存输入数据用于调试
    np.save("debug_input.npy", input_data)
    print("✓ 输入数据已保存至 debug_input.npy")
    
    return input_data

def enhanced_rknn_inference(rknn, input_data):
    """增强的RKNN推理函数"""
    print(f"\n=== RKNN推理执行 ===")
    
    try:
        # 尝试不同的推理配置
        configs = [
            {'inputs': [input_data], 'data_format': ['nchw']},
            {'inputs': [input_data], 'data_format': None},
            {'inputs': [input_data]},
        ]
        
        for i, config in enumerate(configs):
            print(f"尝试配置 {i+1}: {config}")
            try:
                outputs = rknn.inference(**config)
                if outputs and len(outputs) > 0:
                    print(f"✓ 配置 {i+1} 推理成功")
                    break
            except Exception as e:
                print(f"✗ 配置 {i+1} 失败: {str(e)}")
                continue
        else:
            print("所有配置都失败")
            return None
        
        print(f"输出数量: {len(outputs)}")
        
        # 详细检查每个输出
        for i, output in enumerate(outputs):
            print(f"输出 {i}:")
            print(f"  形状: {output.shape}")
            print(f"  数据类型: {output.dtype}")
            print(f"  数值范围: [{output.min():.6f}, {output.max():.6f}]")
            print(f"  统计: 均值={output.mean():.6f}, 标准差={output.std():.6f}")
            print(f"  包含NaN: {np.isnan(output).any()}")
            print(f"  包含Inf: {np.isinf(output).any()}")
            
            # 保存输出用于调试
            np.save(f"debug_rknn_output_{i}.npy", output)
        
        return outputs
        
    except Exception as e:
        print(f"RKNN推理异常: {str(e)}")
        import traceback
        traceback.print_exc()
        return None

def compare_raw_outputs(ultra_output, rknn_output):
    """比较原始推理输出"""
    print(f"\n=== 原始输出比较 ===")
    
    if not rknn_output or len(rknn_output) < 2:
        print("RKNN输出不足，无法比较")
        return
    
    # 检测头比较
    ultra_det = ultra_output[0].numpy()
    rknn_det = rknn_output[0]
    
    print("检测头比较:")
    print(f"  Ultra形状: {ultra_det.shape}")
    print(f"  RKNN形状: {rknn_det.shape}")
    print(f"  Ultra范围: [{ultra_det.min():.6f}, {ultra_det.max():.6f}]")
    print(f"  RKNN范围: [{rknn_det.min():.6f}, {rknn_det.max():.6f}]")
    
    if ultra_det.shape == rknn_det.shape:
        det_diff = np.abs(ultra_det - rknn_det)
        print(f"  最大差异: {det_diff.max():.6f}")
        print(f"  平均差异: {det_diff.mean():.6f}")
        print(f"  差异>0.1的比例: {(det_diff > 0.1).mean()*100:.2f}%")
        
        # 保存差异图
        np.save("debug_detection_diff.npy", det_diff)
    else:
        print("  ✗ 形状不匹配，无法进行数值比较")
    
    # 分割头比较
    ultra_seg = ultra_output[1].numpy()
    rknn_seg = rknn_output[1]
    
    print("分割头比较:")
    print(f"  Ultra形状: {ultra_seg.shape}")
    print(f"  RKNN形状: {rknn_seg.shape}")
    print(f"  Ultra范围: [{ultra_seg.min():.6f}, {ultra_seg.max():.6f}]")
    print(f"  RKNN范围: [{rknn_seg.min():.6f}, {rknn_seg.max():.6f}]")
    
    if ultra_seg.shape == rknn_seg.shape:
        seg_diff = np.abs(ultra_seg - rknn_seg)
        print(f"  最大差异: {seg_diff.max():.6f}")
        print(f"  平均差异: {seg_diff.mean():.6f}")
        print(f"  差异>0.1的比例: {(seg_diff > 0.1).mean()*100:.2f}%")
        
        np.save("debug_segmentation_diff.npy", seg_diff)
    else:
        print("  ✗ 形状不匹配，无法进行数值比较")

def analyze_detection_confidence(ultra_output, rknn_output):
    """分析检测置信度分布"""
    print(f"\n=== 置信度分析 ===")
    
    if not rknn_output or len(rknn_output) < 1:
        print("RKNN输出不足，无法分析")
        return
    
    ultra_det = ultra_output[0].numpy()[0]  # 移除batch维度 [37, 8400]
    rknn_det = rknn_output[0][0]  # 移除batch维度 [37, 8400]
    
    print(f"检测头形状: {ultra_det.shape}")
    
    if ultra_det.shape[0] >= 5:  # 确保有足够的通道
        # 分析objectness (第5个通道，索引4)
        ultra_obj = ultra_det[4, :]
        rknn_obj = rknn_det[4, :]
        
        print("Objectness分析:")
        print(f"  Ultra: 最小={ultra_obj.min():.6f}, 最大={ultra_obj.max():.6f}, 均值={ultra_obj.mean():.6f}")
        print(f"  RKNN: 最小={rknn_obj.min():.6f}, 最大={rknn_obj.max():.6f}, 均值={rknn_obj.mean():.6f}")
        
        # 应用sigmoid激活
        ultra_obj_sigmoid = 1 / (1 + np.exp(-np.clip(ultra_obj, -500, 500)))
        rknn_obj_sigmoid = 1 / (1 + np.exp(-np.clip(rknn_obj, -500, 500)))
        
        print("Objectness (sigmoid后):")
        print(f"  Ultra: 最小={ultra_obj_sigmoid.min():.6f}, 最大={ultra_obj_sigmoid.max():.6f}")
        print(f"  RKNN: 最小={rknn_obj_sigmoid.min():.6f}, 最大={rknn_obj_sigmoid.max():.6f}")
        
        # 统计不同阈值下的激活数量
        thresholds = [0.001, 0.01, 0.05, 0.1, 0.25, 0.5]
        for thresh in thresholds:
            ultra_count = (ultra_obj_sigmoid > thresh).sum()
            rknn_count = (rknn_obj_sigmoid > thresh).sum()
            print(f"  阈值{thresh}: Ultra={ultra_count}, RKNN={rknn_count}")
        
        # 分析类别置信度
        if ultra_det.shape[0] > 5:
            ultra_cls = ultra_det[5:, :]
            rknn_cls = rknn_det[5:, :]
            
            ultra_cls_max = ultra_cls.max(axis=0)
            rknn_cls_max = rknn_cls.max(axis=0)
            
            print("类别置信度分析:")
            print(f"  Ultra最大类别置信度: 最小={ultra_cls_max.min():.6f}, 最大={ultra_cls_max.max():.6f}")
            print(f"  RKNN最大类别置信度: 最小={rknn_cls_max.min():.6f}, 最大={rknn_cls_max.max():.6f}")
            
            # 应用sigmoid到类别预测
            ultra_cls_sigmoid = 1 / (1 + np.exp(-np.clip(ultra_cls_max, -500, 500)))
            rknn_cls_sigmoid = 1 / (1 + np.exp(-np.clip(rknn_cls_max, -500, 500)))
            
            print("类别置信度 (sigmoid后):")
            print(f"  Ultra: 最大={ultra_cls_sigmoid.max():.6f}, 均值={ultra_cls_sigmoid.mean():.6f}")
            print(f"  RKNN: 最大={rknn_cls_sigmoid.max():.6f}, 均值={rknn_cls_sigmoid.mean():.6f}")
            
            # 综合置信度分析
            ultra_combined = ultra_obj_sigmoid * ultra_cls_sigmoid
            rknn_combined = rknn_obj_sigmoid * rknn_cls_sigmoid
            
            print("综合置信度分析:")
            for thresh in thresholds:
                ultra_count = (ultra_combined > thresh).sum()
                rknn_count = (rknn_combined > thresh).sum()
                print(f"  综合阈值{thresh}: Ultra={ultra_count}, RKNN={rknn_count}")

def comprehensive_postprocess_test(predictor, ultra_output, rknn_output, pre_result, img):
    """全面的后处理测试"""
    print(f"\n=== 全面后处理测试 ===")
    
    # 转换RKNN输出为torch张量
    if not rknn_output or len(rknn_output) < 2:
        print("RKNN输出不足")
        return None, None
    
    rknn_det = torch.from_numpy(rknn_output[0])
    rknn_seg = torch.from_numpy(rknn_output[1])
    rknn_infer_result = [rknn_det, rknn_seg]
    
    # 保存原始参数
    original_conf = predictor.args.conf
    original_iou = predictor.args.iou
    
    print(f"原始参数 - 置信度: {original_conf}, IOU: {original_iou}")
    
    # 测试不同的阈值组合
    test_configs = [
        (0.001, 0.45),
        (0.01, 0.45),
        (0.05, 0.45),
        (0.1, 0.45),
        (0.25, 0.45),
        (0.5, 0.45),
        (0.001, 0.1),
        (0.01, 0.1),
        (original_conf, original_iou)
    ]
    
    best_rknn_result = None
    best_config = None
    
    for conf_thresh, iou_thresh in test_configs:
        predictor.args.conf = conf_thresh
        predictor.args.iou = iou_thresh
        
        # Ultra后处理
        ultra_result = predictor.postprocess(ultra_output, pre_result, [img])
        ultra_count = len(ultra_result[0].boxes) if ultra_result and ultra_result[0].boxes is not None else 0
        
        # RKNN后处理
        rknn_result = predictor.postprocess(rknn_infer_result, pre_result, [img])
        rknn_count = len(rknn_result[0].boxes) if rknn_result and rknn_result[0].boxes is not None else 0
        
        print(f"配置(conf={conf_thresh}, iou={iou_thresh}): Ultra={ultra_count}, RKNN={rknn_count}")
        
        # 如果RKNN检测到结果，记录最佳配置
        if rknn_count > 0 and (best_rknn_result is None or rknn_count > len(best_rknn_result[0].boxes)):
            best_rknn_result = rknn_result
            best_config = (conf_thresh, iou_thresh)
            
            # 输出检测到的置信度
            if rknn_result[0].boxes is not None:
                confs = rknn_result[0].boxes.conf.cpu().numpy()
                classes = rknn_result[0].boxes.cls.cpu().numpy()
                print(f"  RKNN检测置信度: {confs}")
                print(f"  RKNN检测类别: {classes}")
    
    # 恢复原始参数
    predictor.args.conf = original_conf
    predictor.args.iou = original_iou
    
    # 使用原始参数进行最终后处理
    final_ultra = predictor.postprocess(ultra_output, pre_result, [img])
    final_rknn = predictor.postprocess(rknn_infer_result, pre_result, [img])
    
    if best_config:
        print(f"\n最佳RKNN配置: conf={best_config[0]}, iou={best_config[1]}")
        print(f"建议使用此配置进行RKNN后处理")
    else:
        print(f"\n所有配置下RKNN都未检测到结果")
        print(f"建议检查:")
        print(f"  1. 模型量化精度")
        print(f"  2. 输入数据预处理")
        print(f"  3. 模型转换参数")
    
    return final_ultra, final_rknn

if __name__ == "__main__":
    # 1. 加载图像
    img_path = "./data/Color/st.jpg"
    img = cv2.imread(img_path)
    if img is None:
        print(f"错误：无法读取图像 {img_path}")
        exit(1)
    
    print(f"图像加载成功，尺寸: {img.shape}")
    original_height, original_width = img.shape[:2]
    
    # 2. 初始化RKNN
    rknn_model_path = './merge_yolo11n_seg_rknn_model/merge_yolo11n_seg-rk3588.rknn'
    rknn = init_rknn_enhanced(rknn_model_path)
    if rknn is None:
        print("RKNN初始化失败")
        exit(1)

    # 3. 加载模型
    seg_model = YOLO("./merge_yolo11n_seg_rknn_model", task="segment")
    seg_model.predict()

    # (1,3,640,640)

    print("✓ Ultralytics模型加载成功")

    # 4. 初始化预测器
    args = dict(
        model="./merge_yolo11n_seg_rknn_model",
        source=img,
        imgsz=640,
        conf=0.25,
        iou=0.45,
        device="cpu"
    )
    predictor = SegmentationPredictor(overrides=args)
    predictor.setup_model(seg_model.model)
    predictor.imgsz = 640
    
    # 5. 设置batch属性
    predictor.batch = [[img_path], [img]]
    
    # 6. 预处理
    pre_result = predictor.preprocess([img])
    print(f"✓ 预处理完成，输入形状: {pre_result.shape}")

    # 7. 验证并准备输入数据
    input_data = validate_and_prepare_input(pre_result)
    if input_data is None:
        print("输入数据准备失败")
        exit(1)

    # 8. Ultralytics推理
    print(f"\n=== Ultralytics推理 ===")
    infer_result = predictor.model(pre_result)
    print(f"✓ Ultralytics推理完成")
    print(f"检测头形状: {infer_result[0].shape}")
    print(f"分割头形状: {infer_result[1].shape}")
    
    # 保存Ultralytics输出
    np.save("debug_ultra_det.npy", infer_result[0].numpy())
    np.save("debug_ultra_seg.npy", infer_result[1].numpy())

    # 9. RKNN推理
    rknn_output = enhanced_rknn_inference(rknn, input_data)
    if rknn_output is None:
        print("RKNN推理失败")
        exit(1)
    
    # 10. 比较原始输出
    compare_raw_outputs(infer_result, rknn_output)
    
    # 11. 分析置信度
    analyze_detection_confidence(infer_result, rknn_output)
    
    # 12. 全面后处理测试
    post_result_ultra, post_result_rknn = comprehensive_postprocess_test(
        predictor, infer_result, rknn_output, pre_result, img
    )
    
    if post_result_ultra:
        ultra_count = len(post_result_ultra[0].boxes) if post_result_ultra[0].boxes is not None else 0
        print(f"\n✓ Ultralytics最终检测到: {ultra_count} 个结果")
    
    if post_result_rknn:
        rknn_count = len(post_result_rknn[0].boxes) if post_result_rknn[0].boxes is not None else 0
        print(f"✓ RKNN最终检测到: {rknn_count} 个结果")
    else:
        print("✗ RKNN未检测到任何结果")
    
    # 13. 可视化并保存结果
    if post_result_ultra:
        print(f"\n保存Ultralytics结果:")
        visualize_and_save_masks(post_result_ultra, img, output_dir="results_ultralytics")
    
    if post_result_rknn:
        print(f"\n保存RKNN结果:")
        visualize_and_save_masks(post_result_rknn, img, output_dir="results_rknn")
    
    print(f"\n=== 调试文件保存位置 ===")
    print(f"输入数据: debug_input.npy")
    print(f"Ultra检测头: debug_ultra_det.npy")
    print(f"Ultra分割头: debug_ultra_seg.npy")
    print(f"RKNN输出: debug_rknn_output_*.npy")
    print(f"差异分析: debug_*_diff.npy")
    
    print(f"\n✓ 所有处理完成")